package com.yupi.springbootinit.guavalimit;

import com.google.common.util.concurrent.RateLimiter;
import com.yupi.springbootinit.common.ErrorCode;
import com.yupi.springbootinit.exception.ThrowUtils;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Aspect
@Component
public class GuavaLimitAop {
    /**
     * 不同的接口，不同的流量控制
     * map的key为 Limiter.key
     */
    private final Map<String, RateLimiter> limitMap = new ConcurrentHashMap<>();

    @Pointcut("execution(* com.yupi.springbootinit.controller.*.*(..)) ")
    public void pc() {
    }


    @Before("pc() && @annotation(guavaLimit)")
    public void before(JoinPoint joinPoint, GuavaLimit guavaLimit) throws Throwable {
        // key作用：不同的接口，不同的流量控制
        String key = guavaLimit.key();
        RateLimiter rateLimiter = null;
        // 验证缓存是否有命中key
        if (!limitMap.containsKey(key)) {
            // 创建令牌桶
            rateLimiter = RateLimiter.create(guavaLimit.permitsPerSecond());
            limitMap.put(key, rateLimiter);
            log.info("新建了令牌桶={}，容量={}", key, guavaLimit.permitsPerSecond());
        }
        rateLimiter = limitMap.get(key);
        // 拿令牌
        boolean acquire = rateLimiter.tryAcquire(guavaLimit.timeout(), guavaLimit.timeunit());
        // 拿不到命令，直接返回异常提示
        ThrowUtils.throwIf(!acquire, ErrorCode.TOO_MANY_REQUEST, guavaLimit.msg());
    }


}
